#!/usr/bin/env python3
"""
手势分类模型训练脚本
使用MobileNetV2迁移学习构建静态手势分类模型
"""

import json
import os

import matplotlib.pyplot as plt
import numpy as np
import seaborn as sns
from sklearn.metrics import classification_report, confusion_matrix
from tensorflow.keras.applications import MobileNetV2
from tensorflow.keras.callbacks import ModelCheckpoint, EarlyStopping, ReduceLROnPlateau
from tensorflow.keras.layers import GlobalAveragePooling2D, Dense, Dropout
from tensorflow.keras.models import Model
from tensorflow.keras.optimizers import Adam

# 配置matplotlib中文字体
plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei', 'SimSun']  # 优先使用黑体、微软雅黑、宋体
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

def build_model(num_classes, input_shape=(224, 224, 3)):
    """
    构建基于MobileNetV2的迁移学习模型
    """
    print("正在构建模型...")
    
    # 加载预训练的MobileNetV2模型（不包括顶层）
    base_model = MobileNetV2(
        weights='imagenet',
        include_top=False,
        input_shape=input_shape
    )
    
    # 冻结基础模型的权重
    base_model.trainable = False
    
    # 添加自定义顶层
    inputs = base_model.input
    x = base_model.output
    
    # 全局平均池化
    x = GlobalAveragePooling2D()(x)
    
    # 全连接层
    x = Dense(128, activation='relu')(x)
    x = Dropout(0.3)(x)
    
    # 输出层
    outputs = Dense(num_classes, activation='softmax')(x)
    
    # 创建模型
    model = Model(inputs=inputs, outputs=outputs)
    
    print("模型结构：")
    model.summary()
    
    return model

def compile_model(model, learning_rate=0.001):
    """
    编译模型
    """
    print("正在编译模型...")
    
    # 使用Adam优化器
    optimizer = Adam(learning_rate=learning_rate)
    
    # 编译模型
    model.compile(
        optimizer=optimizer,
        loss='categorical_crossentropy',
        metrics=['accuracy']
    )
    
    return model

def train_model(model, train_generator, val_generator, epochs=10):
    """
    训练模型
    """
    print("开始训练模型...")
    
    # 创建回调函数
    callbacks = [
        # 保存最佳模型到models目录
        ModelCheckpoint(
            'models/static_gesture_classifier_best.keras',
            monitor='val_accuracy',
            save_best_only=True,
            mode='max',
            verbose=1
        ),
        # 早停
        EarlyStopping(
            monitor='val_loss',
            patience=5,
            restore_best_weights=True,
            verbose=1
        ),
        # 学习率调整
        ReduceLROnPlateau(
            monitor='val_loss',
            factor=0.2,
            patience=3,
            min_lr=1e-7,
            verbose=1
        )
    ]
    
    # 训练模型
    history = model.fit(
        train_generator,
        epochs=epochs,
        validation_data=val_generator,
        callbacks=callbacks,
        verbose=1
    )
    
    return history

def fine_tune_model(model, train_generator, val_generator, epochs=5):
    """
    微调模型（解冻部分层）
    """
    print("开始微调模型...")
    
    # 找到MobileNetV2基础模型
    base_model = None
    for layer in model.layers:
        if 'mobilenet' in layer.name.lower():
            base_model = layer
            break
    
    if base_model is None:
        print("警告：未找到MobileNetV2基础模型，跳过微调")
        return None
    
    # 解冻基础模型
    base_model.trainable = True
    
    # 设置微调策略：只解冻最后几层
    # 找到基础模型的层数
    base_layers = base_model.layers
    num_layers = len(base_layers)
    
    # 冻结前80%的层，解冻后20%的层
    fine_tune_at = int(num_layers * 0.8)
    
    for layer in base_layers[:fine_tune_at]:
        layer.trainable = False
    
    for layer in base_layers[fine_tune_at:]:
        layer.trainable = True
    
    print(f"微调设置：解冻了 {num_layers - fine_tune_at} 层（总共 {num_layers} 层）")
    
    # 重新编译模型（使用较低的学习率）
    model.compile(
        optimizer=Adam(learning_rate=0.0001),
        loss='categorical_crossentropy',
        metrics=['accuracy']
    )
    
    # 微调训练
    history_fine = model.fit(
        train_generator,
        epochs=epochs,
        validation_data=val_generator,
        verbose=1
    )
    
    return history_fine

def evaluate_model(model, test_generator, class_names):
    """
    评估模型性能
    """
    print("正在评估模型...")
    
    # 预测测试集
    y_pred = model.predict(test_generator)
    y_pred_classes = np.argmax(y_pred, axis=1)
    
    # 获取真实标签 - 从生成器中提取
    # 由于我们使用的是ImageDataGenerator，需要手动获取真实标签
    y_true = []
    for i in range(len(test_generator)):
        _, batch_labels = test_generator[i]
        batch_true = np.argmax(batch_labels, axis=1)
        y_true.extend(batch_true)
    
    y_true = np.array(y_true)
    
    # 分类报告
    print("\n分类报告：")
    print(classification_report(y_true, y_pred_classes, target_names=class_names))
    
    # 混淆矩阵
    cm = confusion_matrix(y_true, y_pred_classes)
    
    plt.figure(figsize=(10, 8))
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', 
                xticklabels=class_names, yticklabels=class_names)
    plt.title('混淆矩阵')
    plt.xlabel('预测标签')
    plt.ylabel('真实标签')
    plt.tight_layout()
    # 保存到docs目录
    os.makedirs('docs', exist_ok=True)
    plt.savefig('docs/confusion_matrix.png', dpi=300, bbox_inches='tight')
    plt.show()
    
    # 计算准确率
    accuracy = np.sum(y_true == y_pred_classes) / len(y_true)
    print(f"测试集准确率: {accuracy:.4f}")
    
    return accuracy

def plot_training_history(history, history_fine=None):
    """
    绘制训练历史
    """
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 5))
    
    # 绘制准确率
    ax1.plot(history.history['accuracy'], label='训练准确率')
    ax1.plot(history.history['val_accuracy'], label='验证准确率')
    
    if history_fine:
        ax1.plot(np.arange(len(history.history['accuracy']), 
                         len(history.history['accuracy']) + len(history_fine.history['accuracy'])),
                history_fine.history['accuracy'], label='微调训练准确率')
        ax1.plot(np.arange(len(history.history['val_accuracy']), 
                         len(history.history['val_accuracy']) + len(history_fine.history['val_accuracy'])),
                history_fine.history['val_accuracy'], label='微调验证准确率')
    
    ax1.set_title('模型准确率')
    ax1.set_xlabel('周期')
    ax1.set_ylabel('准确率')
    ax1.legend()
    ax1.grid(True)
    
    # 绘制损失
    ax2.plot(history.history['loss'], label='训练损失')
    ax2.plot(history.history['val_loss'], label='验证损失')
    
    if history_fine:
        ax2.plot(np.arange(len(history.history['loss']), 
                         len(history.history['loss']) + len(history_fine.history['loss'])),
                history_fine.history['loss'], label='微调训练损失')
        ax2.plot(np.arange(len(history.history['val_loss']), 
                         len(history.history['val_loss']) + len(history_fine.history['val_loss'])),
                history_fine.history['val_loss'], label='微调验证损失')
    
    ax2.set_title('模型损失')
    ax2.set_xlabel('周期')
    ax2.set_ylabel('损失')
    ax2.legend()
    ax2.grid(True)
    
    plt.tight_layout()
    # 保存到docs目录
    os.makedirs('docs', exist_ok=True)
    plt.savefig('docs/training_history.png', dpi=300, bbox_inches='tight')
    plt.show()

def load_data_generators():
    """
    加载数据生成器（从原始位置读取，节省内存）
    """
    try:
        # 尝试导入数据预处理模块的函数
        from preprocess_data import create_dataframe_from_directory, create_data_generators_from_dataframe
        
        # 扫描数据目录
        df, class_names = create_dataframe_from_directory("data/leapGestRecog")
        
        if df is None:
            print("错误：数据扫描失败")
            return None, None, None, None
        
        # 创建数据生成器（从原始位置读取）
        train_generator, val_generator, test_generator, label_to_index = create_data_generators_from_dataframe(
            df,
            class_names,
            batch_size=32,
            validation_split=0.2
        )
        
        if train_generator is None:
            print("错误：数据生成器创建失败")
            return None, None, None, None
        
        data_size = train_generator.samples + val_generator.samples
        
        return train_generator, val_generator, class_names, data_size
        
    except ImportError as e:
        print(f"导入错误: {e}")
        print("请确保preprocess_data.py文件存在")
        return None, None, None, None
    except Exception as e:
        print(f"数据加载错误: {e}")
        import traceback
        traceback.print_exc()
        return None, None, None, None

def main():
    """主函数"""
    print("=" * 60)
    print("手势识别项目 - 模型训练模块")
    print("=" * 60)
    
    # 检查数据预处理是否完成
    if not os.path.exists('models/class_mapping.json'):
        print("错误：请先运行数据预处理脚本")
        print("命令: python preprocess_data.py")
        return
    
    # 加载类别映射
    with open('models/class_mapping.json', 'r') as f:
        label_map = json.load(f)
    
    class_names = list(label_map.keys())
    num_classes = len(class_names)
    
    print(f"手势类别数量: {num_classes}")
    print(f"类别名称: {class_names}")
    
    # 加载数据生成器
    print("正在加载数据生成器...")
    train_generator, val_generator, loaded_class_names, data_size = load_data_generators()
    
    if train_generator is None:
        print("数据生成器加载失败，程序退出")
        return
    
    print(f"数据加载完成：共 {data_size} 个样本")
    
    # 构建模型
    model = build_model(num_classes)
    
    # 编译模型
    model = compile_model(model)
    
    # 训练模型（使用完整训练轮数）
    print("开始模型训练...")
    history = train_model(model, train_generator, val_generator, epochs=20)
    
    # 跳过微调阶段，直接评估
    print("跳过微调阶段，直接评估模型...")
    history_fine = None
    
    # 评估模型
    print("开始模型评估...")
    accuracy = evaluate_model(model, val_generator, class_names)
    
    # 绘制训练历史
    plot_training_history(history, history_fine)
    
    # 保存最终模型到models目录
    model.save('models/static_gesture_classifier_final.keras')
    
    print("\n" + "=" * 60)
    print("✓ 模型训练完成！")
    print("=" * 60)
    print(f"最终测试准确率: {accuracy:.4f}")
    print("\n生成的文件：")
    print("- models/static_gesture_classifier_best.keras: 最佳验证准确率模型")
    print("- models/static_gesture_classifier_final.keras: 最终训练模型")
    print("- docs/training_history.png: 训练历史图表")
    print("- docs/confusion_matrix.png: 混淆矩阵")
    print("\n模型训练完成！")

if __name__ == "__main__":
    main()